import os
import shutil
from sklearn.model_selection import train_test_split

# 数据集路径
data_dir = '../dataset'  # 修改为正确的相对路径
# 目标路径
output_dir = 'prepared_dataset'

# 创建训练和验证目录
train_dir = os.path.join(output_dir, 'train')
val_dir = os.path.join(output_dir, 'val')
os.makedirs(train_dir, exist_ok=True)
os.makedirs(val_dir, exist_ok=True)

# 获取所有类别
categories = os.listdir(data_dir)

for category in categories:
    category_path = os.path.join(data_dir, category)
    images = os.listdir(category_path)
    train_images, val_images = train_test_split(images, test_size=0.2, random_state=42)

    # 创建类别目录
    train_category_dir = os.path.join(train_dir, category)
    val_category_dir = os.path.join(val_dir, category)
    os.makedirs(train_category_dir, exist_ok=True)
    os.makedirs(val_category_dir, exist_ok=True)

    # 复制图像到训练集和验证集目录
    for image in train_images:
        shutil.copy(os.path.join(category_path, image), os.path.join(train_category_dir, image))

    for image in val_images:
        shutil.copy(os.path.join(category_path, image), os.path.join(val_category_dir, image))

print("数据集准备完成")
